#include "pch.h"
#include "rtspmedia.h"


using namespace std::chrono;

uint64_t GetTickCount64CrossPlatform() {
	using namespace std::chrono;
	return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
}

RtspMedia::RtspMedia()
	:video_codec_id(MKCodecH264),
	audio_codec_id(MKCodecAAC)
{
	audio_buffer = std::make_unique<char[]>(48000 * 2);
	memset(audio_buffer.get(), 0, 48000 * 2);
}

RtspMedia::~RtspMedia()
{
}

void RtspMedia::StartCapture()
{
	_desktop_video = std::make_unique<DesktopCapture>(0, this);
	if (_desktop_video->Init())
	{
		if (_src_w==0|| _src_h==0)
		{
			_desktop_video->GetVideoInfo(_src_w, _src_h, _src_fps);
		}
		if (video_codec_id == MKCodecVP8)
		{
			_video_encoder = std::make_unique<VP8Encoder>(_src_w, _src_h, 25, 8096, this);
		}
		else if (video_codec_id == MKCodecH264)
		{
			_video_encoder = std::make_unique<X264Encoder>(_src_w, _src_h, 25, 8096, this);
		}
		_desktop_video->Start();
	}

	CLoopbackCapture::GetInstance().SetCallBack(shared_from_this());
	CLoopbackCapture::GetInstance().Start();
	CLoopbackCapture::GetInstance().GetAudioInfo(_channels, _sampels);

	if (audio_codec_id == MKCodecOpus)
	{
		audio_frame_size = 480;
		_audio_encoder = std::make_unique<OpusEncode>(_sampels, _channels, 128, this);
	}
	else if (audio_codec_id == MKCodecAAC)
	{
		audio_frame_size = 1024;
		_audio_encoder = std::make_unique<AACEncoder>(_sampels, _channels, 128, this);
	}
}

void RtspMedia::StopCapture()
{
	if (_desktop_video)
	{
		_desktop_video->Stop();
		_desktop_video.reset();
	}
	CLoopbackCapture::GetInstance().Stop();
}

void RtspMedia::StartPush()
{
	_media = mk_media_create("__defaultVhost__", "live", "desktop", 0, 0, 0);	
	codec_args video_codec_args;
	video_codec_args.video.width = _src_w;
	video_codec_args.video.height = _src_h;
	video_codec_args.video.fps = 25;
	_video_track = mk_track_create(video_codec_id, &video_codec_args);

	codec_args audio_codec_args;
	audio_codec_args.audio.channels = _channels;
	audio_codec_args.audio.sample_rate = _sampels;
	_audio_track = mk_track_create(audio_codec_id, &audio_codec_args);
	mk_media_init_track(_media, _video_track);
	mk_media_init_track(_media, _audio_track);
	_start_time = GetTickCount64CrossPlatform();

	_isStartPush = true;
}

void RtspMedia::StopPush()
{
	_isStartPush = false;
	_src_w = _src_h = 0;
	mk_media_release(_media);
}

void RtspMedia::onVideoCaptureData(void* data, int len)
{
	if (!_isStartPush)
	{
		return;
	}
	if (_video_encoder)
	{
		long long pts = GetTickCount64CrossPlatform() - _start_time;
		_video_encoder->encode(data,len, pts);
	}
}

void RtspMedia::onAudioCaptureData(void* data, int len)
{
	if (!_isStartPush)
	{
		return;
	}
	if (_audio_encoder)
	{		
		memcpy(audio_buffer.get() + audio_samples * 2, data, len);
		audio_samples += len/ _channels /2;
		if (audio_samples >= audio_frame_size)
		{
			_audio_encoder->encode(audio_buffer.get(), audio_frame_size * 2,0);
			audio_samples -= audio_frame_size;
			memmove(audio_buffer.get(), audio_buffer.get() + audio_frame_size * 2, audio_samples * 2);
		}
	}
}

void RtspMedia::onVideoEncodeData(void* data, long long len, int dts, int pts)
{
	mk_frame frame = mk_frame_create(video_codec_id, dts, pts, (const char*)data, len, nullptr, nullptr);
	mk_media_input_frame(_media, frame);
	mk_frame_unref(frame);
}

void RtspMedia::onAudioEncodeData(void* data, long long len)
{
	long long pts = GetTickCount64CrossPlatform() - _start_time;
	mk_frame frame = mk_frame_create(audio_codec_id, pts, pts, (const char*)data, len, nullptr, nullptr);
	mk_media_input_frame(_media, frame);
	mk_frame_unref(frame);
}


